import pkg_resources
pkg_resources.require("matplotlib")
import os
import sys
import re
import pickle
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as tk
import math
#import pyplot
#from pylab import *
FONTSIZE=16
from optparse import OptionParser

def read_curve(path, filenames):
	datapoint={}
	for root, dirs, files in os.walk(path):
		if (len(files) > 0):
		        pt = []
		        for filename in filenames:
		                fn = os.path.join(root,filename)
		                try:
		                        f=open(fn)
		                except IOError:
		                        print "NOT FOUND:" + fn
		                        continue
		                pt.append(float(f.readline()))
		                f.close()
		        if(len(pt)>0):
				ntasks=[s for s in os.path.basename(root).split('-') if s.isdigit()]
				if len(ntasks)!=1:
					print "Number of tasks is not detected"
		                #ntasks = re.sub('[SMPVN]','',os.path.basename(root))
		                #print root + ":\t"+ntasks[0]+" "+ str(max(pt))
		                datapoint[int(ntasks[0])]=max(pt)/0.2
				print int(ntasks[0]), datapoint[int(ntasks[0])]
	return datapoint
 
def dict_to_x_and_y(data):
    x = []
    y = []
 
    #  ensure the data is sorted by date, i.e. dict key
    sorted_keys = data.keys()
    sorted_keys.sort()
    for key in sorted_keys:
        #  parse date to system date object
        x.append(int(key))
        y.append(float(data[key]))
 
    return (x, y)
def format_xaxis(x, pos=None):
	if x%1024 != 0:
		res = str(int(x))
	else:
		res = str(int(x/1024))+'K'
	return res
def format_yaxis(y, pos=None):
	if y<1.0:
		res = str(y)
	else:
		res = str(int(y))
	return res

parser = OptionParser()

parser.add_option("-p", "--path", dest="path", default=".",
                  help="specify the path where to search for datafiles")
parser.add_option("-n", "--filenames", dest="filenames", default="events1 events2", help=" names of the files that contains datpoints. Use space to separate. \"event1 event2\" is default;")
parser.add_option("-l", "--label", dest="label", default="", help=" label for the curves. Use semicolon as a separator. Directory path is the default value.")
parser.add_option("-w", "--weak", dest="weak", action ="store_true",default=False,
                  help="weak scaling diagram;")
parser.add_option("-a", "--axis", dest="axis", default="64,32768,0.5,128", help=" axis limits. Use comma to separate. <xmin,xmax, ymin,ymax > = <64,32768,0.5,128> is default")

(options, args) = parser.parse_args()
filenames  = options.filenames.lower().split(' ')
labels = options.label.split(';')
[xmin,xmax, ymin, ymax] = [float(val) for val in options.axis.split(',')]
ttl = "Strong"
if(options.weak):
	ttl = "Weak"
#plt.title(ttl+ " scaling performance on JUQUEEN",fontsize=FONTSIZE)
xtitle = "total # of processors [K=1024]"
ytitle = "run-time/ biological time"
plt.ylabel(ytitle,fontsize=FONTSIZE)
plt.xlabel(xtitle,fontsize=FONTSIZE)
plt.xticks(fontsize=FONTSIZE)
plt.yticks(fontsize=FONTSIZE)
plt.xscale('log', basex=2)
plt.yscale('log', basey=2)
plt.axis([xmin,xmax, ymin, ymax])
current_figure = plt.gcf() 
current_axes = plt.gca()
current_axes.xaxis.set_major_formatter(tk.FuncFormatter(format_xaxis))
current_axes.yaxis.set_major_formatter(tk.FuncFormatter(format_yaxis))
current_axes.set_autoscale_on(False)
plt.grid()
c=0
color=['k','b','g','r', 'y']
style=['^','o','s','*']
filled=['black', 'white']
x_max = []
base_folder=''
for root, dirs, files in os.walk(options.path):
	if (len(dirs)==0 ): 
		print root
		parent_folder=root.split('/')[-2];
		if ( len(base_folder)>0 and base_folder==parent_folder):
			continue;
		else:
			base_folder=parent_folder
		#print root[0:root.rindex('/')];
		datapoints = read_curve(root[0:root.rindex('/')], filenames)
		(x, y) = dict_to_x_and_y(datapoints)
		if len(x) > len(x_max):
			x_max = x
		lbl = root[0:root.rindex('/')]
		if(len(labels) > 1):
			lbl = labels[c]
		plt.plot(x, y, color[0]+style[c/2], linestyle='-',  markerfacecolor=filled[c%2],linewidth=0.5, markersize=7, label=lbl)
		c+=1
		

#plot strong scaling line
if(options.weak):
	y=[ymax/4]*(len(x_max))
else:
	y=[(ymax*2)/i for i in [math.pow(2,i) for i in xrange(1,len(x_max)+1)]]
#plt.plot(x_max,y,color[0], linestyle='--', linewidth=0.5, label='')
lgd = plt.legend( loc=1, numpoints = 1 )
lgd.set_visible(True)

plt.show()







